"""
检索服务
实现基于向量相似度的文本检索
"""
import numpy as np
from typing import List, Dict, Any, Tuple
from app.storage.json_storage import KnowledgeBaseStorage
from app.services.embedding import embedding_service
from app.core.config import settings


class RetrievalService:
    """检索服务类"""
    
    def __init__(self, storage: KnowledgeBaseStorage):
        self.storage = storage
    
    def retrieve(self, query: str, top_k: int = None) -> Tuple[List[str], List[float], List[str]]:
        """
        检索最相关的文本块
        
        Args:
            query: 查询文本
            top_k: 返回的文本块数量，默认使用配置值
            
        Returns:
            (文本列表, 相似度分数列表, 来源文档列表)
        """
        if top_k is None:
            top_k = settings.TOP_K
        
        # 获取所有文本块
        chunks = self.storage.get_all_chunks()
        if not chunks:
            return [], [], []
        
        # 生成查询向量
        query_embedding = embedding_service.create_embedding(query)
        query_vector = np.array(query_embedding)
        
        # 计算相似度
        similarities = []
        for chunk in chunks:
            chunk_vector = np.array(chunk["embedding"])
            # 使用点积计算余弦相似度（假设向量已归一化）
            similarity = np.dot(chunk_vector, query_vector)
            similarities.append(similarity)
        
        # 获取 top_k 个最相似的索引
        similarities = np.array(similarities)
        top_indices = np.argsort(similarities)[-top_k:][::-1]
        
        # 提取结果
        texts = [chunks[i]["text"] for i in top_indices]
        scores = [float(similarities[i]) for i in top_indices]
        documents = [chunks[i]["document"] for i in top_indices]
        
        return texts, scores, documents


def create_retrieval_service() -> RetrievalService:
    """创建检索服务实例"""
    storage = KnowledgeBaseStorage(
        storage_path=settings.knowledge_base_path,
        embedding_model=settings.EMBEDDING_MODEL
    )
    return RetrievalService(storage)

